"""
RAG问答引擎
"""
from typing import List, Dict, Optional, AsyncGenerator
from app.services.llm_client import ollama_client
from app.services.vector_store import VectorStore
import logging

logger = logging.getLogger(__name__)


class RAGEngine:
    """RAG（检索增强生成）问答引擎"""
    
    def __init__(self, vector_store: Optional[VectorStore] = None):
        """
        初始化RAG引擎
        
        Args:
            vector_store: 向量存储实例
        """
        self.llm_client = ollama_client
        self.vector_store = vector_store
    
    async def answer_question(
        self,
        question: str,
        context_docs: Optional[List[str]] = None,
        temperature: float = 0.7,
        max_tokens: int = 2000,
    ) -> str:
        """
        回答问题（非流式）
        
        Args:
            question: 用户问题
            context_docs: 上下文文档列表
            temperature: 温度参数
            max_tokens: 最大token数
            
        Returns:
            回答内容
        """
        try:
            # 构建系统提示
            system_prompt = self._build_system_prompt()
            
            # 如果有上下文文档，添加到提示中
            if context_docs:
                context = "\n\n".join([f"文档片段 {i+1}:\n{doc}" for i, doc in enumerate(context_docs)])
                user_prompt = f"""参考以下文档内容回答问题：

{context}

问题：{question}

请基于上述文档内容给出详细、准确的回答。如果文档中没有相关信息，请明确说明。"""
            else:
                user_prompt = question
            
            # 构建消息
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]
            
            # 调用LLM
            answer = await self.llm_client.chat_completion(
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens
            )
            
            return answer
        except Exception as e:
            logger.error(f"问答失败: {str(e)}")
            raise
    
    async def answer_question_stream(
        self,
        question: str,
        context_docs: Optional[List[str]] = None,
        temperature: float = 0.7,
        max_tokens: int = 2000,
    ) -> AsyncGenerator[str, None]:
        """
        回答问题（流式输出）
        
        Args:
            question: 用户问题
            context_docs: 上下文文档列表
            temperature: 温度参数
            max_tokens: 最大token数
            
        Yields:
            回答的文本片段
        """
        try:
            system_prompt = self._build_system_prompt()
            
            if context_docs:
                context = "\n\n".join([f"文档片段 {i+1}:\n{doc}" for i, doc in enumerate(context_docs)])
                user_prompt = f"""参考以下文档内容回答问题：

{context}

问题：{question}

请基于上述文档内容给出详细、准确的回答。如果文档中没有相关信息，请明确说明。"""
            else:
                user_prompt = question
            
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ]
            
            async for chunk in self.llm_client.chat_completion_stream(
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens
            ):
                yield chunk
        except Exception as e:
            logger.error(f"流式问答失败: {str(e)}")
            raise
    
    async def retrieve_and_answer(
        self,
        question: str,
        top_k: int = 5,
        temperature: float = 0.7,
        max_tokens: int = 2000,
        stream: bool = False
    ):
        """
        检索相关文档并回答问题
        
        Args:
            question: 用户问题
            top_k: 检索文档数量
            temperature: 温度参数
            max_tokens: 最大token数
            stream: 是否流式输出
            
        Returns:
            回答内容或生成器
        """
        if not self.vector_store or self.vector_store.count() == 0:
            # 没有向量库，直接问答
            if stream:
                return self.answer_question_stream(
                    question=question,
                    temperature=temperature,
                    max_tokens=max_tokens
                )
            else:
                return await self.answer_question(
                    question=question,
                    temperature=temperature,
                    max_tokens=max_tokens
                )
        
        try:
            # 生成问题的向量
            question_embedding = await self.llm_client.generate_embedding(question)
            
            # 检索相关文档
            results = self.vector_store.similarity_search(
                query_embedding=question_embedding,
                k=top_k
            )
            
            # 提取文档内容（不过滤，因为FAISS已经返回top_k最相关的）
            context_docs = [text for text, score, metadata in results]
            
            logger.info(f"检索到 {len(context_docs)} 个相关文档片段（总共{len(results)}个结果）")
            for i, (text, score, metadata) in enumerate(results[:3]):
                logger.info(f"  文档{i+1} 相似度:{score:.4f} 内容:{text[:50]}...")
            
            # 回答问题
            if stream:
                return self.answer_question_stream(
                    question=question,
                    context_docs=context_docs,
                    temperature=temperature,
                    max_tokens=max_tokens
                )
            else:
                return await self.answer_question(
                    question=question,
                    context_docs=context_docs,
                    temperature=temperature,
                    max_tokens=max_tokens
                )
        except Exception as e:
            logger.error(f"检索问答失败: {str(e)}")
            raise
    
    def _build_system_prompt(self) -> str:
        """构建系统提示"""
        return """你是一个专业的AI助手，专注于回答用户的问题。

请遵循以下原则：
1. 基于提供的文档内容准确回答问题
2. 如果文档中没有相关信息，诚实地告诉用户
3. 回答要清晰、详细、有条理
4. 使用简洁的中文表达
5. 如果问题不清楚，可以要求用户澄清

请开始回答用户的问题。"""


# 全局RAG引擎实例（可以根据知识库动态创建）
def create_rag_engine(vector_store: Optional[VectorStore] = None) -> RAGEngine:
    """创建RAG引擎实例"""
    return RAGEngine(vector_store=vector_store)

